import json
import logging
import time
import uuid  # 用于生成唯一 Request ID
from datetime import datetime  # 用于精确时间戳
from typing import List

from fastapi import Request, Response
from rich.console import Console
from starlette.concurrency import iterate_in_threadpool
from starlette.middleware.base import BaseHTTPMiddleware

logger = logging.getLogger(__name__)
console = Console()


class LoggingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next) -> Response:
        # 排除 Socket.IO 请求（polling/WebSocket 握手），避免日志噪声
        if request.url.path.startswith("/socket.io"):
            return await call_next(request)  # 直接传递，不记录日志

        # 生成唯一 Request ID，用于关联请求和响应
        request_id = str(uuid.uuid4())

        # 记录开始时间
        start_time = time.perf_counter()

        # 缓存请求体（Bytes），允许下游复用
        body_bytes = await request.body()

        # 解析请求参数
        query_params = dict(request.query_params)
        path_params = request.path_params
        try:
            request_body = json.loads(body_bytes) if body_bytes else {}
        except json.JSONDecodeError:
            request_body = body_bytes.decode("utf-8") if body_bytes else {}

        request_log = {
            "request_id": request_id,  # 新增：唯一 ID
            "method": request.method,
            "url": request.url.path,
            "query_params": query_params,
            "path_params": path_params,
            "request_body": request_body,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[
                :-3
            ],  # 精确到毫秒
        }
        console.print("[bold green]🚀 请求开始[/bold green]", style="bold cyan")
        console.print_json(data=request_log, indent=2, highlight=True)

        logger.info(
            f"Request ID: {request_id} - Method: {request.method} - URL: {request.url.path}"
        )

        # 调用下一个处理链（下游可复用 Body）
        response = await call_next(request)

        # 记录结束时间
        end_time = time.perf_counter()
        duration = (end_time - start_time) * 1000  # 毫秒

        # 收集响应体 chunks（仅一次迭代）
        raw_chunks: List[bytes] = [chunk async for chunk in response.body_iterator]

        # 重置响应迭代器，允许下游/客户端复用
        response.body_iterator = iterate_in_threadpool(iter(raw_chunks))

        # 解析响应体
        try:
            response_body = (
                json.loads(raw_chunks[0].decode("utf-8")) if raw_chunks else {}
            )
        except (json.JSONDecodeError, IndexError, UnicodeDecodeError):
            response_body = "非 JSON 响应"

        response_log = {
            "request_id": request_id,  # 同一 ID，用于关联
            "url": request.url.path,
            "status_code": response.status_code,
            "response_body": response_body,
            "duration_ms": round(duration, 2),
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[
                :-3
            ],  # 精确到毫秒
        }
        console.print("[bold yellow]📤 响应完成[/bold yellow]", style="bold magenta")
        console.print_json(data=response_log, indent=2, highlight=True)
        if duration > 500:  # 如果响应慢于 500ms，突出警告
            console.print(f"[bold red]⚠️  响应较慢: {duration:.2f}ms[/bold red]")

        logger.info(
            f"Response ID: {request_id} - Status: {response.status_code} - Duration: {duration:.2f}ms"
        )

        return response
